%% DESCRIPTION

% This script generates stimulus-triggered averages of cortical evoked responses.
% You can analyze a single file or combine two files together to geneate the average responses.
% Individual neural sweeps are baseline-corrected by subtracting the mean of the PRE_SWEEP
% from all the samples in the sweep.
% The script generates plots of the stimulus triggers (Figure 1), individual torque sweeps (for a 
% specified channel; Figure 2), corrected torque sweeps (for the same channel; Figure 3) and the average 
% evoked and corrected evoked responses for the specified torque channels (Figures 4 & 5).

%% DEFINE VARIABLES
baseFilenames = [
    "20180215_006"   
];

for n = 1 : length(baseFilenames)

    fprintf('Processing %d of %d files\n', n, length(baseFilenames));

    baseFilename = baseFilenames{n};
    BASEDIR = '/Users/samiramoorjani/Documents/Ubi/';
    DATADIR = [BASEDIR 'Data/'];
    figureOutputPathPrefix = [BASEDIR 'Figures_BCofSweeps/' baseFilename '_torque_']; % This is where the figures will be stored
    ns4Filename = [baseFilename '.ns4'];

    fprintf('%s\n', ns4Filename);
    
    % Specify PRE and POST in milliseconds
    PRE = 20;
    POST = 160;

    AMPLIFICATION = 6553.4; % Max. digital value / Max. analog value for analog channels (V)
    SCALAR = 0.5; % Torque scaling factor

    % Enter experiment information
    TORQUE = 'F-E Torque';
    CURRENT = 160; % Enter stimulation current in �A 
    SITE = 'LM1(14,15)I/D'; % Enter site of stimulation
    
    % Select 'Combine Files' or 'Load Files.' If using 'Combine Files,' specify the two filenames. 
    % Verify Outliers 
    % Verify list of channels to be plotted
    % Verify Y-axis limits
    % Verify parameters for 'rmoutliers_sd' function

    %% COMBINE FILES

    % Specify the two files that need to be combined together
    % ns4Filename1 = '20180215_001.ns4';
    % ns4Filename2 = '20180215_001.ns5';
    % 
    % % Combine the two files together into a single NSx structure in MATLAB
    % [NS4] = combineNSx([DATADIR ns4Filename1], [DATADIR ns4Filename2]);

    %% LOAD FILES

    % If the .ns4 file has already been loaded, don't re-load the file
    if exist('filenameExtracted', 'var') ~= 1 || strcmp(filenameExtracted, "") || ~strcmp(filenameExtracted, ns4Filename)
        fprintf('Loading NS4 file: %s\n', ns4Filename);
        NS4 = openNSx('report', 'read', [DATADIR ns4Filename], 'p:double');
        if exist('NS4', 'var') ~= 1 || isstruct(NS4) ~= 1
            filenameExtracted = "";
            return;
        end
        filenameExtracted = ns4Filename;
    else
        fprintf('Loading already-loaded NS4 file: %s\n', filenameExtracted);
    end

    %% EXTRACT STIMULUS TRIGGERS

    fprintf('Extracting stimulus-trigger locations\n');

    % Specify location of stimulus triggers
    triggerData = NS4.Data(end, :);

    % Baseline correction
    triggerData = triggerData - mean(triggerData); 

    % Stop underscore from being subcripted in the printed filenames
    underscoreStartIndex = strfind(ns4Filename, '_');
    prefix_ns4Filename = ns4Filename(1 : underscoreStartIndex - 1);
    postfix_ns4Filename = ns4Filename(underscoreStartIndex + 1 : end);
    printNs4Filename = [prefix_ns4Filename '\_' postfix_ns4Filename];

    % Find local maxima
    [ peaks, triggerLocations ] = findpeaks(double(triggerData));

    % Remove noise peaks
    PEAK_THRESHOLD = 1000;
    triggerLocations = triggerLocations(peaks > PEAK_THRESHOLD);

    % Specify sampling frequency
    samplingFrequency = NS4.MetaTags.SamplingFreq;

    % Convert PRE and POST to samples
    PRE_SWEEP = fix(PRE * samplingFrequency / 1000);
    POST_SWEEP = fix(POST * samplingFrequency / 1000);

    sweepLength = PRE_SWEEP + POST_SWEEP; 

    % Remove triggers that are too close together
    for t = 1 : length(triggerLocations) - 1

        if triggerLocations(t + 1) - triggerLocations(t) < sweepLength
           triggerLocations(t) = NaN;
%            triggerLocations(t + 1) = NaN;
        end

    end

    % Remove first and last triggers due to missing PRE_SWEEP or POST_SWEEP
    % triggerLocations(1) = NaN;
    triggerLocations(length(triggerLocations)) = NaN;

    % Remove outliers
%     triggerLocations(11 : 18) = NaN;
%     triggerLocations(27) = NaN;

    correctedTriggerLocations = triggerLocations(not(isnan(triggerLocations)));

    % Plot triggerData
    figure
    plot(triggerData);
    title(sprintf('Graph of file ''%s''; Total number of triggers: %d, Corrected number of triggers: %d', printNs4Filename, length(triggerLocations), length(correctedTriggerLocations)))
    grid on
    
    drawnow
    savefig([figureOutputPathPrefix '1.fig'])

    %% GENERATE TORQUE SWEEPS

    % Specify location of torque channels 
    channelLocations = [ 129 ];
    channelNames = {NS4.ElectrodesInfo(channelLocations).Label};

    triggerCount = length(correctedTriggerLocations);
    channelCount = length(channelLocations);   

    sweeps = nan(sweepLength, triggerCount, channelCount);
    correctedSweeps = nan(sweepLength, triggerCount, channelCount);
    
    fprintf('Processing data\n');

    for c = 1 : channelCount 
        channelData = NS4.Data(channelLocations(c), :);   

        for t = 1 : triggerCount  
            sweepStart = correctedTriggerLocations(t) - PRE_SWEEP;
            sweepEnd = sweepStart + sweepLength - 1;
            sweep = channelData(sweepStart : sweepEnd);

            % Baseline correction of individual sweeps by subtracting the mean of the first 150 samples in the PRE_SWEEP
            baselineCorrectionFactor = mean(sweep(1, 1 : 150), 2); 
            sweep = sweep(1, :) - baselineCorrectionFactor; 
            
            sweeps(:, t, c) = sweep;    
        end   
        
        correctedSweeps = rmoutliers_sd(sweeps, 2, 2.7, 50); % Removes outlier sweeps. Format: rmoutliers_sd(data, dimension, sd, nmin)
    
    end
    
    % Plot sweeps for a specified channel
    channelId = 1;

    % Stop underscore from being subcripted in the printed channel names
    underscoreStartIndex = strfind(channelNames{channelId}, '_');
    prefixChannelName = channelNames{channelId}(1 : underscoreStartIndex - 1);
    postfixChannelName = channelNames{channelId}(underscoreStartIndex + 1 : end);
    printChannelName = [prefixChannelName '\_' postfixChannelName];
    
    figure
    plot(sweeps(:, :, channelId))
    ylim([-10000  10000])
    title(sprintf('Graph of file ''%s''; Total number of sweeps: %d;  Site ''%s''', printNs4Filename, size(sweeps, 2), printChannelName))
    grid on
    
    drawnow
    savefig([figureOutputPathPrefix '2.fig'])
   
    % Plot corrected sweeps for the same channel
    figure
    plot(correctedSweeps(:, :, channelId))
    ylim([-10000  10000])
    title(sprintf('Graph of file ''%s''; Corrected number of sweeps: %d;  Site ''%s''', printNs4Filename, size(correctedSweeps, 2), printChannelName))
    grid on
    
    drawnow
    savefig([figureOutputPathPrefix '3.fig'])

    %% PLOT EVOKED RESPONSES
    
    % Plot evoked responses from all sweeps
    mEPBefore = mean(sweeps, 2); % mean(A,2) returns a column vector containing the mean of the elements in each row
    mEP = squeeze(mEPBefore); % Removes singleton dimensions to make it a 2D matrix with data for each channel appearing in a separate column

    fprintf('Generating plots\n');

    % Specify torque channel 
    torqueChannelId = 1; % 1 is F-E torque
     
    % Plot torque channel 
    timeAxis = 1000 / samplingFrequency : 1000 / samplingFrequency : 1000 * (size(mEP,1)) / samplingFrequency; % X-axis is in milliseconds
    timeAxis = timeAxis - PRE; % Forces t = 0 at trigger location
    singleChannelMEP = mEP(:, torqueChannelId); 
    
    figure
    plot(timeAxis, (singleChannelMEP / AMPLIFICATION) * SCALAR, 'Color', [1, 0.5, 0], 'LineWidth', 1.5', 'Displayname', sprintf('''%s''', printNs4Filename))
    ylim([-0.05 0.05])
    xlabel({'Time (ms)'}, 'FontSize', 12, 'FontWeight', 'bold')
    ylabel({'Amplitude (N.m)'}, 'FontSize', 12, 'FontWeight', 'bold')
    title(sprintf('Recorded ''%s'' at %0.1f �A delivered to ''%s''', TORQUE, CURRENT, SITE))
    legend('show')
    
    drawnow
    savefig([figureOutputPathPrefix '4.fig'])

    % Plot evoked responses from corrected sweeps
    correctedMEPBefore = mean(correctedSweeps, 2); % mean(A,2) returns a column vector containing the mean of the elements in each row
    correctedMEP = squeeze(correctedMEPBefore); % Removes singleton dimensions to make it a 2D matrix with data for each channel appearing in a separate column
     
    % Plot torque channel 
    timeAxis = 1000 / samplingFrequency : 1000 / samplingFrequency : 1000 * (size(mEP,1)) / samplingFrequency; % X-axis is in milliseconds
    timeAxis = timeAxis - PRE; % Forces t = 0 at trigger location
    singleChannelMEP = correctedMEP(:, torqueChannelId); 
    
    figure
    plot(timeAxis, (singleChannelMEP / AMPLIFICATION) * SCALAR, 'Color', [1, 0.5, 0], 'LineWidth', 1.5', 'Displayname', sprintf('''%s''', printNs4Filename))
    ylim([-0.05 0.05])
    xlabel({'Time (ms)'}, 'FontSize', 12, 'FontWeight', 'bold')
    ylabel({'Amplitude (N.m)'}, 'FontSize', 12, 'FontWeight', 'bold')
    title(sprintf('Recorded corrected ''%s'' at %0.1f �A delivered to ''%s''', TORQUE, CURRENT, SITE))
    legend('show')
    
    drawnow
    savefig([figureOutputPathPrefix '5.fig'])
    
end
